{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eec680f5",
   "metadata": {},
   "source": [
    "# LangGraph 활용 Agent 구축\n",
    "\n",
    "이번 튜토리얼에서는 웹 검색 도구를 통해 챗봇에 웹 검색 기능수행하는 Agent 을 추가합니다.\n",
    "\n",
    "LLM 에 도구를 바인딩하여 LLM 에 입력된 요청에 따라 필요시 웹 검색 도구(Tool)를 호출하는 Agent 을 구축합니다.\n",
    "\n",
    "뿐만아니라, 조건부 엣지를 통해 도구 호출 여부에 따라 다른 노드로 라우팅하는 방법도 함께 배워봅니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de9d9d8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5c6228",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Modules\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d5376bf",
   "metadata": {},
   "source": [
    "## 도구(Tool) 사용하기\n",
    "\n",
    "**참고**\n",
    "\n",
    "- [도구(Tools)](https://wikidocs.net/262582)\n",
    "\n",
    "챗봇이 \"기억\"에서 답변할 수 없는 질문을 처리하기 위해 웹 검색 도구를 통합할 것입니다. 이 도구를 사용하여 관련 정보를 찾아 더 나은 응답을 제공할 수 있습니다.\n",
    "\n",
    "### 검색 API 도구\n",
    "\n",
    "Tavily 검색 API를 활용하여 검색 기능을 구현하는 도구입니다. 이 도구는 두 가지 주요 클래스를 제공합니다: `TavilySearchResults`와 `TavilyAnswer`.\n",
    "\n",
    "**API 키 발급 주소**\n",
    "- https://app.tavily.com/\n",
    "\n",
    "발급한 API 키를 환경변수에 설정합니다.\n",
    "\n",
    "`.env` 파일에 아래와 같이 설정합니다.\n",
    "\n",
    "```\n",
    "TAVILY_API_KEY=tvly-abcdefghijklmnopqrstuvwxyz\n",
    "```\n",
    "\n",
    "### TavilySearchResults\n",
    "\n",
    "**설명**\n",
    "- Tavily 검색 API를 쿼리하고 JSON 형식의 결과를 반환합니다.\n",
    "- 포괄적이고 정확하며 신뢰할 수 있는 결과에 최적화된 검색 엔진입니다.\n",
    "- 현재 이벤트에 대한 질문에 답변할 때 유용합니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81cbcea8",
   "metadata": {},
   "source": [
    "다음으로 웹 검색 도구인 `TavilySearchResults`를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "163da255",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -U langchain-teddynote"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04fccdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.tools.tavily import TavilySearch\n",
    "\n",
    "# 검색 도구 생성\n",
    "tool = TavilySearch(max_results=3)\n",
    "\n",
    "# 도구 목록에 추가\n",
    "tools = [tool]\n",
    "\n",
    "# 도구 실행\n",
    "print(tool.invoke(\"테디노트 랭체인 튜토리얼\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b51b6ca3",
   "metadata": {},
   "source": [
    "결과는 챗봇이 질문에 답할 수 있도록 사용할 수 있는 페이지 요약입니다.\n",
    "\n",
    "이번에는 LLM에 `bind_tools`를 추가하여 **LLM + 도구** 를 구성합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d6166da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph.message import add_messages\n",
    "\n",
    "\n",
    "# State 정의\n",
    "class State(TypedDict):\n",
    "    # list 타입에 add_messages 적용(list 에 message 추가)\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efe9e3c7",
   "metadata": {},
   "source": [
    "LLM 을 정의하고 도구를 바인딩합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8c65ea60",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# LLM 에 도구 바인딩\n",
    "llm_with_tools = llm.bind_tools(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc0c2841",
   "metadata": {},
   "source": [
    "노드를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "028d36e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 노드 함수 정의\n",
    "def chatbot(state: State):\n",
    "    answer = llm_with_tools.invoke(state[\"messages\"])\n",
    "    # 메시지 목록 반환\n",
    "    return {\"messages\": [answer]}  # 자동으로 add_messages 적용"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c798e0",
   "metadata": {},
   "source": [
    "그래프 생성 및 노드를 추가합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e8d16a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph\n",
    "\n",
    "# 상태 그래프 초기화\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 노드 추가\n",
    "graph_builder.add_node(\"chatbot\", chatbot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a034ee74",
   "metadata": {},
   "source": [
    "## 도구 노드(Tool Node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89b08e3f",
   "metadata": {},
   "source": [
    "다음으로, 도구가 호출될 경우 실제로 실행할 수 있는 함수를 만들어야 합니다. 이를 위해 새로운 노드에 도구를 추가합니다.\n",
    "\n",
    "가장 최근의 메시지를 확인하고 메시지에 `tool_calls`가 포함되어 있으면 도구를 호출하는 `BasicToolNode`를 구현합니다. \n",
    "\n",
    "지금은 직접 구현하지만, 나중에는 LangGraph의 pre-built 되어있는 [ToolNode](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.ToolNode) 로 대체할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1437765",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from langchain_core.messages import ToolMessage\n",
    "\n",
    "\n",
    "class BasicToolNode:\n",
    "    \"\"\"Run tools requested in the last AIMessage node\"\"\"\n",
    "\n",
    "    def __init__(self, tools: list) -> None:\n",
    "        # 도구 리스트\n",
    "        self.tools_list = {tool.name: tool for tool in tools}\n",
    "\n",
    "    def __call__(self, inputs: dict):\n",
    "        # 메시지가 존재할 경우 가장 최근 메시지 1개 추출\n",
    "        if messages := inputs.get(\"messages\", []):\n",
    "            message = messages[-1]\n",
    "        else:\n",
    "            raise ValueError(\"No message found in input\")\n",
    "\n",
    "        # 도구 호출 결과\n",
    "        outputs = []\n",
    "        for tool_call in message.tool_calls:\n",
    "            # 도구 호출 후 결과 저장\n",
    "            tool_result = self.tools_list[tool_call[\"name\"]].invoke(tool_call[\"args\"])\n",
    "            outputs.append(\n",
    "                # 도구 호출 결과를 메시지로 저장\n",
    "                ToolMessage(\n",
    "                    content=json.dumps(\n",
    "                        tool_result, ensure_ascii=False\n",
    "                    ),  # 도구 호출 결과를 문자열로 변환\n",
    "                    name=tool_call[\"name\"],\n",
    "                    tool_call_id=tool_call[\"id\"],\n",
    "                )\n",
    "            )\n",
    "\n",
    "        return {\"messages\": outputs}\n",
    "\n",
    "\n",
    "# 도구 노드 생성\n",
    "tool_node = BasicToolNode(tools=[tool])\n",
    "\n",
    "# 그래프에 도구 노드 추가\n",
    "graph_builder.add_node(\"tools\", tool_node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de52ecda",
   "metadata": {},
   "source": [
    "## 조건부 엣지(Conditional Edge)\n",
    "\n",
    "도구 노드가 추가되면 `conditional_edges`를 정의할 수 있습니다.\n",
    "\n",
    "**Edges**는 한 노드에서 다음 노드로 제어 흐름을 라우팅합니다. \n",
    "\n",
    "**Conditional edges**는 일반적으로 \"if\" 문을 포함하여 현재 그래프 상태에 따라 다른 노드로 라우팅합니다. 이러한 함수는 현재 그래프 `state`를 받아 다음에 호출할 Node 를 나타내는 **문자열 또는 문자열 목록** 을 반환합니다.\n",
    "\n",
    "아래에서는 `route_tools`라는 라우터 함수를 정의하여 챗봇의 출력에서 `tool_calls`를 확인합니다. \n",
    "\n",
    "이 함수를 `add_conditional_edges`를 호출하여 그래프에 제공하면, `chatbot` 노드가 완료될 때마다 이 함수를 확인하여 다음으로 어디로 갈지 결정합니다.\n",
    "\n",
    "조건은 도구 호출이 있으면 `tools`로, 없으면 `END`로 라우팅됩니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- langgraph 에 pre-built 되어 있는 [tools_condition](https://langchain-ai.github.io/langgraph/reference/prebuilt/#tools_condition) 으로 대체할 수 있습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f758bf4f",
   "metadata": {},
   "source": [
    "### `add_conditional_edges`\n",
    "\n",
    "![add_conditional_edges](./image/langgraph-02.png)\n",
    "\n",
    "`add_conditional_edges` 메서드는 시작 노드에서 여러 대상 노드로의 조건부 엣지를 추가합니다.\n",
    "\n",
    "**매개변수**\n",
    "- `source` (str): 시작 노드. 이 노드를 나갈 때 조건부 엣지가 실행됩니다.\n",
    "- `path` (Union[Callable, Runnable]): 다음 노드를 결정하는 호출 가능한 객체 또는 Runnable. `path_map`을 지정하지 않으면 하나 이상의 노드를 반환해야 합니다. `END`를 반환하면 그래프 실행이 중지됩니다.\n",
    "- `path_map` (Optional[Union[dict[Hashable, str], list[str]]]): 경로와 노드 이름 간의 매핑. 생략하면 `path`가 반환하는 값이 노드 이름이어야 합니다.\n",
    "- `then` (Optional[str]): `path`로 선택된 노드 실행 후 실행할 노드의 이름.\n",
    "\n",
    "**반환값**\n",
    "- Self: 메서드 체이닝을 위해 자기 자신을 반환합니다.\n",
    "\n",
    "**주요 기능**\n",
    "1. 조건부 엣지를 그래프에 추가합니다.\n",
    "2. `path_map`을 딕셔너리로 변환합니다.\n",
    "3. `path` 함수의 반환 타입을 분석하여 자동으로 `path_map`을 생성할 수 있습니다.\n",
    "4. 조건부 분기를 그래프에 저장합니다.\n",
    "\n",
    "**참고**\n",
    "- 이미 컴파일된 그래프에 엣지를 추가하면 경고 메시지가 출력됩니다.\n",
    "- `path` 함수의 반환 값에 대한 타입 힌트가 없거나 `path_map`이 제공되지 않으면, 그래프 시각화 시 해당 엣지가 그래프의 모든 노드로 전환될 수 있다고 가정합니다.\n",
    "- 동일한 이름의 분기가 이미 존재하는 경우 `ValueError`가 발생합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6a964c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import START, END\n",
    "\n",
    "\n",
    "def route_tools(\n",
    "    state: State,\n",
    "):\n",
    "    if messages := state.get(\"messages\", []):\n",
    "        # 가장 최근 AI 메시지 추출\n",
    "        ai_message = messages[-1]\n",
    "    else:\n",
    "        # 입력 상태에 메시지가 없는 경우 예외 발생\n",
    "        raise ValueError(f\"No messages found in input state to tool_edge: {state}\")\n",
    "\n",
    "    # AI 메시지에 도구 호출이 있는 경우 \"tools\" 반환\n",
    "    if hasattr(ai_message, \"tool_calls\") and len(ai_message.tool_calls) > 0:\n",
    "        # 도구 호출이 있는 경우 \"tools\" 반환\n",
    "        return \"tools\"\n",
    "    # 도구 호출이 없는 경우 \"END\" 반환\n",
    "    return END\n",
    "\n",
    "\n",
    "# `tools_condition` 함수는 챗봇이 도구 사용을 요청하면 \"tools\"를 반환하고, 직접 응답이 가능한 경우 \"END\"를 반환\n",
    "graph_builder.add_conditional_edges(\n",
    "    source=\"chatbot\",\n",
    "    path=route_tools,\n",
    "    # route_tools 의 반환값이 \"tools\" 인 경우 \"tools\" 노드로, 그렇지 않으면 END 노드로 라우팅\n",
    "    path_map={\"tools\": \"tools\", END: END},\n",
    ")\n",
    "\n",
    "# tools > chatbot\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "# START > chatbot\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# 그래프 컴파일\n",
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beab199c",
   "metadata": {},
   "source": [
    "**조건부 엣지**가 단일 노드에서 시작해야 합니다.\n",
    "\n",
    "이는 그래프에 \"`chatbot`\" 노드가 실행될 때마다 도구를 호출하면 'tools'로 이동하고, 직접 응답하면 루프를 종료하라는 의미입니다. \n",
    "\n",
    "사전 구축된 `tools_condition`처럼, 함수는 도구 호출이 없을 경우 `END` 문자열을 반환(그래프 종료) 합니다. 그래프가 `END`로 전환되면 더 이상 완료할 작업이 없으며 실행을 중지합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d4d1118",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b736b473",
   "metadata": {},
   "source": [
    "이제 봇에게 훈련 데이터 외의 질문을 할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31343b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\"messages\": \"테디노트 YouTube 채널에 대해서 검색해 줘\"}\n",
    "\n",
    "for event in graph.stream(inputs, stream_mode=\"values\"):\n",
    "    for key, value in event.items():\n",
    "        print(f\"\\n==============\\nSTEP: {key}\\n==============\\n\")\n",
    "        # display_message_tree(value[\"messages\"][-1])\n",
    "        print(value[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4c6c5ac",
   "metadata": {},
   "source": [
    "도구 호출 후 구조에 대한 이미지"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23623eec",
   "metadata": {},
   "source": [
    "![](./image/tool-message-01.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff003852",
   "metadata": {},
   "source": [
    "- [이전 실행에 대한 LangSmith 추적](https://smith.langchain.com/public/4f82ddfa-a452-40f3-ab09-4eb088b812a4/r)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
